# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jaewon


class _DAG_Node:
  def __init__(self, node_name, node_idx):
    self.name = node_name
    self.idx = node_idx
    self.in_nodes = set()
    self.out_nodes = set()

  def add_in_node(self, node):
    self.in_nodes.add(node)
    node.out_nodes.add(self)

  def freeze(self):
    self.in_nodes = tuple(self.in_nodes)
    self.out_nodes = tuple(self.out_nodes)

  def __repr__(self):
    return '_DAG_Node(\'%s\', %d)' % (self.name, self.idx)


class _DAG:
  def __init__(self, graph):
    self._nodes = []
    self._node_map = {}

    for idx, n in enumerate(graph.nodes):
      dnode = _DAG_Node(n.name, idx)
      self._nodes.append(dnode)
      self._node_map[n.name] = dnode

      for arg_name, var in n.inputs.items():
        dep_node = self._node_map[var.node.name]
        dnode.add_in_node(dep_node)

    for dnode in self._nodes:
      dnode.freeze()

  def query_transitive_deps(self, node_name):
    visited = [False] * len(self._nodes)
    transitive_deps = set()

    def visit_in_nodes(node):
      visited[node.idx] = True
      for d in node.in_nodes:
        if visited[d.idx]:
          continue
        transitive_deps.add(d)
        visit_in_nodes(d)

    visit_in_nodes(self._node_map[node_name])
    return sorted(transitive_deps, key=lambda n: n.idx)

  def query_transitive_rdeps(self, node_name):
    visited = [False] * len(self._nodes)
    transitive_rdeps = set()

    def visit_out_nodes(node):
      visited[node.idx] = True
      for d in node.out_nodes:
        if visited[d.idx]:
          continue
        transitive_rdeps.add(d)
        visit_out_nodes(d)

    visit_out_nodes(self._node_map[node_name])
    return sorted(transitive_rdeps, key=lambda n: n.idx)


def build_graph_queryer(graph):
  return _DAG(graph)
